import pandas as pd
from nlpx.model import TextCNN
from nlpx.model.wrapper import FastClassifyModelWrapper
from nlpx.llm import AlbertTokenizeVec, ErnieTokenizeVec, BertTokenizeVec

# sentence-transformers/all-MiniLM-L6-v2
pretrained_path = r'/Users/summy/data/nlp_structbert_backbone_base_std'


if __name__ == '__main__':
	# from modelscope import AutoTokenizer

	df = pd.read_csv('~/project/python/parttime/text_gcn/data/北方地区不安全事件统计20240331.csv', encoding='GBK')
	df['故障标志'] = df['故障标志'].astype('category')
	classes = df['故障标志'].cat.categories.values.tolist()
	df['故障标志'] = df['故障标志'].cat.codes
	y = df['故障标志'].to_numpy()

	tokenize_vec = BertTokenizeVec(pretrained_path)

	text_cnn = TextCNN(tokenize_vec.hidden_size, out_features=len(classes))
	model = FastClassifyModelWrapper(tokenize_vec, classes=classes)
	model.train(text_cnn, df['故障描述'].values.tolist(), y, n_jobs=3, num_workers=2)

	result = model.predict_classes_proba(['昆明航B737飞机执行昆明至西安航班 西安机场进近过程中  机组将飞机襟翼位置设置错误触发近地警告事件',
						   				  '海航 B737飞机执行汉中至广州航班 飞机起飞后发现汉中过站期间货舱有730KG货物未卸机'])
	print(result)
	# tokenizer = AutoTokenizer.from_pretrained(pretrained_path, revision='v1.0.0')
	# print(tokenizer.(df['故障描述'].values.tolist()))
